"""
Author:cold
Date:2021-04-06
Version:10.0
Info: ExtraTreesRegressor
"""
from pandas import read_csv
from sklearn.model_selection import train_test_split,StratifiedKFold,cross_val_score
from sklearn.ensemble import ExtraTreesRegressor
from sklearn.metrics import mean_squared_error
from keras.wrappers.scikit_learn import KerasRegressor
import pandas as pd



# 加载数据(455)
dataset =read_csv('train_dataset.csv').values


# 划分训练集和测试集
X = dataset[:,0:13]
Y = dataset[:,13]

x_train,x_test,y_train,y_test = train_test_split(X,Y,test_size=0.3,random_state=7)


# 创建线性回归模型
# lr = LinearRegression()
abr =  ExtraTreesRegressor()
model =abr
# 拟合训练数据

estimator = KerasRegressor(build_fn=model ,nb_epoch=100, batch_size=3, verbose=1)
kfold = StratifiedKFold(n_splits=10, shuffle=True, random_state=7)
cross_val_score(estimator, X, Y, cv=kfold)
model.fit(x_train,y_train)
# 得到预测结果
y_test_pred = model.predict(x_test)
y_train_pred = model.predict(x_train)


# 计算相应的评测指标
error_test = mean_squared_error(y_test,y_test_pred)
error_train = mean_squared_error(y_train,y_train_pred)
print("训练集误差为：{}，测试集误差为：{}".format(error_train,error_test))


#预测结果
testset =read_csv('test_dataset.csv').values
x_pred = testset[:,1:14]
y_pred = model.predict(x_pred)
ID = []
for i in range(len(y_pred)):
    ID.append("id_"+str(i+1))
res = pd.DataFrame()
res['ID']=ID
res['value']=y_pred
res.to_csv('res.csv',index=False)
print("res.csv 已生成")